from trimesh.exchange.binvox import load_binvox
import trimesh
import numpy as np
import airsim
from airsim.types import Vector3r
import os


class AirsimPLYGenerator:
    def __init__(
        self,
        client: airsim.client.MultirotorClient,
        ridius=2,
        min_x=-3,
        max_x=5,
        min_y=-2.5,
        max_y=2.5,
        min_z=0,
        max_z=4,
        resolution=0.02,
    ):
        self.client = client
        self.ridius = ridius
        self.min_x = min_x
        self.max_x = max_x
        self.min_y = min_y
        self.max_y = max_y
        self.min_z = min_z
        self.max_z = max_z
        self.x_length = int((max_x - min_x) / ridius + 0.5)
        self.y_length = int((max_y - min_y) / ridius + 0.5)
        self.z_length = int((max_z - min_z) / ridius + 0.5)
        self.resolution = resolution
        

    def _binvox_to_pc(self, binvox_file, start_x, start_y, start_z):
        half_ridius = self.ridius / 2
        ridius = self.ridius
        if not os.path.exists(binvox_file):
            return False, np.array([])
        with open(binvox_file, "rb") as f:
            voxel_grid = load_binvox(f)
        try:
            points = (voxel_grid.points + half_ridius) * ridius * ridius - half_ridius
            x = points[:, 1] + half_ridius + start_x
            y = -points[:, 0] + start_y
            z = points[:, 2] + start_z
            points = np.stack([x, y, z], axis=1)
            return True, points
        except:
            return False, np.array([])

    def _create_ptclouds(self, start_x, start_y, start_z, binvox_name):
        half_ridius = self.ridius / 2
        ridius = self.ridius
        pos = Vector3r(half_ridius + start_x, -start_y, -start_z)
        try:
            self.client.simCreateVoxelGrid(
                pos, ridius, ridius, ridius, self.resolution, binvox_name
            )
        except:
            print("Error")

    def generate(self, path):
        print("Start to generate point clouds, it may take a while.")
        points_list = []
        tmp_path = os.path.abspath(path)
        tmp_path = os.path.join(tmp_path, "tmp")
        if not os.path.exists(tmp_path):
            os.makedirs(tmp_path)
        total_process = self.x_length * self.y_length * self.z_length
        for i in range(self.x_length):
            for j in range(self.y_length):
                for k in range(self.z_length):
                    start_x = i * self.ridius + self.min_x
                    start_y = j * self.ridius + self.min_y
                    start_z = k * self.ridius + self.min_z
                    binvox_name = os.path.join(tmp_path, f"tmp_{i}_{j}_{k}.binvox")
                    self._create_ptclouds(start_x, start_y, start_z, binvox_name)
                    success, points = self._binvox_to_pc(
                        binvox_name, start_x, start_y, start_z
                    )
                    if success:
                        points_list.append(points)
                    process = (
                        i * self.y_length * self.z_length + j * self.z_length + k
                    ) / total_process
                    yield process
        if len(points_list) == 0:
            print("No point clouds generated.")
        else:
            points_list = np.concatenate(points_list)
            # Create trimesh point cloud
            point_cloud = trimesh.points.PointCloud(points_list)
            # Export to .ply
            output_path = os.path.join(path, "model.ply")
            point_cloud.export(output_path)
            print(f"Exported to {os.path.abspath(output_path)}")
            print("Start to reconstruct the mesh")
        for i in range(self.x_length):
            for j in range(self.y_length):
                for k in range(self.z_length):
                    binvox_name = os.path.join(tmp_path, f"tmp_{i}_{j}_{k}.binvox")
                    os.remove(binvox_name)
        os.rmdir(tmp_path)
        yield 1.0


def __main__():
    client = airsim.MultirotorClient("192.168.0.100")
    client.confirmConnection()
    client.simFlushPersistentMarkers()
    ply_generator = AirsimPLYGenerator(client)
    for process in ply_generator.generate("."):
        print(f"Process: {process * 100:.2f}%")
        
if __name__ == "__main__":
    __main__()
